package org.taj.common.tcp;

import org.taj.common.tcp.exception.ConnectionException;
import org.taj.common.tcp.exception.MessageException;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.zeromq.SocketType;
import org.zeromq.ZContext;
import org.zeromq.ZMQ;
import org.zeromq.ZMQException;

import java.io.Closeable;
import java.io.IOException;

/**
 * 客户端
 */
public class TcpClient implements Closeable {
    private static final Logger logger = LogManager.getLogger(TcpClient.class);
    private final String connIP;
    private final int connPort;
    private final int subPort;
    private final ZContext ctx;
    private ZMQ.Socket connSocket;
    private ZMQ.Socket subSocket;
    private ResponseCallback cb;
    private ZMQ.Poller poller;
    private final String appName;

    public TcpClient(String ip, int port, String appName) {
        this.connIP = ip;
        this.connPort = port;
        this.subPort = port + 1;
        ctx = new ZContext();
        poller = ctx.createPoller(1);
        this.appName = appName;
    }

    public boolean connect() {
        try {
            connSocket = ctx.createSocket(SocketType.DEALER);
            boolean b = connSocket.connect("tcp://" + connIP + ":" + connPort);
            subSocket = ctx.createSocket(SocketType.SUB);
            b &= subSocket.connect("tcp://" + connIP + ":" + subPort);
            poller = ctx.createPoller(2);
            poller.register(connSocket, ZMQ.Poller.POLLIN);
            poller.register(subSocket, ZMQ.Poller.POLLIN);
            new Thread(() -> {
                while (!Thread.currentThread().isInterrupted()) {
                    poller.poll();
                    if(poller.pollin(0)) {
                        byte[] data = connSocket.recv();
                        if (cb != null) {
                            cb.recv("-", data);
                        }
                    }
                    if (poller.pollin(1)) {
                        String topic = subSocket.recvStr(); // 接收主题
                        byte[] data = recvBytes();
                        if (cb != null) {
                            cb.recv(topic, data);
                        }
                    }
                }
            }).start();
            if(b) {
                b = connSocket.send(appName);
            }
            return b;
        } catch (ZMQException ex) {
            throw new ConnectionException(ex);
        }
    }

    /**
     * 发送消息
     */
    public boolean sendBytes(byte[] data) {
        Message.Messenger.Builder msgBuilder = Message.Messenger.newBuilder();
        Message.Messenger messenger = msgBuilder
                .setHeader(Message.Header.newBuilder().setCode(1).build())
                .setError(Message.Error.newBuilder().setCode(0))
                .setBody(ByteString.copyFrom(data)).build();
        return connSocket.send(messenger.toByteArray());
    }


    /**
     * 元数据接受
     */
    private byte[] recvRawBytes() {
        byte[] res = connSocket.recv();
        Message.Messenger.Builder msgBuilder = Message.Messenger.newBuilder();
//        if (res.length == 1) {
//            Message.Messenger messenger = msgBuilder
//                    .setHeader(Message.Header.newBuilder()
//                            .setCode(0)
//                            .setBuf0(ByteString.copyFrom("single byte message", StandardCharsets.UTF_8))
//                            .build())
//                    .setBody(ByteString.copyFrom(res)).build();
//            return messenger.toByteArray();
//        }
        return res;
    }

    /**
     * 消息接受
     */
    public byte[] recvBytes() {
        byte[] res = connSocket.recv();
        try {
            Message.Messenger messenger = Message.Messenger.parseFrom(res);
            if (messenger.hasHeader()) {
                // TODO: 处理消息头
            }
            if (messenger.hasError()) {
                // TODO: 处理错误
            }
            return messenger.getBody().toByteArray();
        } catch (InvalidProtocolBufferException ex) {
            throw new MessageException(ex);
        }
    }

    public byte[] sendBytesWaitAck(byte[] data) {
        if (this.sendBytes(data)) {
            return recvBytes();
        }
        return new byte[0];
    }

    public boolean subscribe(String topic) {
        return subSocket.subscribe(topic);
    }

    public boolean unsubscribe(String topic) {
        return subSocket.unsubscribe(topic);
    }

    public void setTopicCallback(ResponseCallback cb) {
        this.cb = cb;
    }

    @Override
    public void close() throws IOException {
        logger.info("客户端主动关闭");
        connSocket.send(new byte[0]);
        connSocket.close();
        subSocket.close();
        ctx.close();
    }
}
